/*
  The scheduler.c file implements the scheduler.
*/

#include <types.h>
#include <klib.h>
#include <scheduler.h>
#include <a.out.h>
#include <linkedlist.h>
#include <i386_arch.h>
#include <i386_pmode.h>
#include <panic.h>
#include <pit.h>
#include <linkedlist.h>
#include <vmm.h>
#include <spinlock.h>

// This is used in the scheduler to keep releasing a periodic job
// is it time to release a job?
// if it's period is a multiple of the current time minus the job's phase
// ct = current time
// p  = job
//#define RELEASEABLE(ct,p)   ((((ct) - (p->phase)) % (p->period)) == 0)


// List of executable process images
// it is initialized by "stuffer"
volatile ptr_list_t process_image_list=NULL;

volatile pid_t latest_pid=1;

// lists of processes
// pending for acceptance
volatile ptr_list_t ap_process_list=NULL;
// real time processes
volatile ptr_list_t rt_process_list=NULL;
// round robin processes (non real-time)
volatile ptr_list_t rr_process_list=NULL;

// pointer to current process
volatile process_t *current_process=NULL;

spinlock_t scheduler_mutex;

double load_factor;

// internal prototypes
process_memory_t *initialize_memory_map(image_t *exec_image);
void release_memory_map(process_memory_t *map);
int calculate_num_pages(dword_t byte_size);
image_t *find_image(char *name);
void process_initialize_registers(process_t *proc);

void make_idle_task(void);
void make_init_task(void);
void system_idle_process(void);

// functions implementation


void initialize_scheduler(void)
{
  process_t *p;
  int freq;

  // initialize mutex
  lock_initialize(&scheduler_mutex);

  load_factor = 0.0;

  // initialize lists of processes (empty)
  ap_process_list=llist_make_list();
  rt_process_list=llist_make_list();
  rr_process_list=llist_make_list();


  // create the IDLE process
  make_idle_task();

  // put the IDLE process into the RR list
  p=llist_get_first(ap_process_list);
  llist_add_last(rr_process_list,p);

  
  // create the INIT process
  make_init_task();

  // put the INIT process into the RR list
  p=llist_get_first(ap_process_list);
  llist_add_first(rr_process_list,p);

  // prepare the PIT
  // number of micro-seconds in a second divided by TIME_QUANTUM 
  // (TIME_QUANTUM is expressed in micro-seconds)
  freq = 1000000 / TIME_QUANTUM;   
  pit_set_frequency(freq);
  assign_IRQ(0,(dword_t)&pit_isr);

  // multitasking will take place as soon as the scheduler interrupt (PIT) 
  // or another interrupt is fired
}

void make_idle_task(void)
{
  process_t *p;
  uregs_t *stack_regs;

  // create init process (system process)
  p=create_process(NULL);
  // set it ready
  p->state=PS_READY;

  // give it a name
  strcpy(p->name,"idle");
  
  // use this address space
  p->cr3=read_cr3();

  // set the entry point
  stack_regs=(uregs_t *)(p->kstack_top);
  stack_regs->eip=(dword_t)system_idle_process;

  // set the current process to it
  current_process=p;
}

void system_idle_process(void)
{
  for (;;)
    {
      //asm("sti");
      asm("hlt");
    }
}


void make_init_task(void)
{
  process_t *p;
  image_t *first_img;

  // get first image to run
  first_img=llist_peek_first(process_image_list);
  // create init process
  p=create_process(first_img);
  // set it ready
  p->state=PS_READY;

  // set the current process to it
  //current_process=p;

  // switch to its address space
  //write_cr3(p->cr3);
}

void schedule(void)
{
  //ptr_iterator_t iter;
  ptr_node_t pnode;
  process_t *p,*chosen;
  qword_t current_time;
  qword_t earliest_deadline;
  int hwi;

  // acquire mutex
  hwi=lock_acquire_cli(&scheduler_mutex);

  // set current_process to ready if it was running
  if (current_process->state==PS_RUNNING)
    {
      current_process->state=PS_READY; 

      // if it's a normal process (Round-Robin), refill it's time slice
      if (current_process->type==TYPE_NORMAL && current_process->time_slice <= 0)
	current_process->time_slice=TIME_SLICE;
    } 

  chosen=NULL;
  earliest_deadline=0;

  //pit_get_time_ms(&current_time);

  // Look for a Real-Time process to run
  // according to EDF criterion
  for (pnode=rt_process_list->header->next; pnode!=rt_process_list->header; pnode=pnode->next)
    {
      p=(process_t *)(pnode->data);

      /*
      // release periodic jobs if applicable
      // or remove non-periodic dead ones
      if (p->state == PS_TERMINATED)
	{
	  if ((p->type == TYPE_PERIODIC) && RELEASEABLE(current_time,p))
	    {
	      p->exec_remaining = p->exec_time;
	      p->absolute_deadline = current_time + p->deadline;
	      process_initialize_registers(p);
	      p->state = PS_READY;
	    }
	}
      */

      // if we haven't chosen a process yet, choose one...
      if ((earliest_deadline==0) && (p->state == PS_READY))
	{
	  earliest_deadline=p->absolute_deadline;
	  chosen=p;
	  continue;
	}

      // if P's deadline is earlier... then choose P
      if ((p->state == PS_READY) && (p->absolute_deadline < earliest_deadline))
	{
	  earliest_deadline=p->absolute_deadline;
	  chosen=p;
	}
    }

  // if we haven't found a Real-Time process to run, 
  // look for in the Round-Robin list
  // at least we will find the IDLE process which is always ready  
  if (chosen == NULL)
    do {
      // move it to the end of the queue (Round-Robin)
      chosen=llist_get_first(rr_process_list);
      llist_add_last(rr_process_list,chosen);
    } while (chosen->state != PS_READY);

  /*
  // this should never happen
  if (chosen==NULL)
    panic("No process to run!\n");
  */
  
  //printf("SCHED: going to execute: %s\n",chosen->name);

  // set the chosen process to be running
  chosen->state=PS_RUNNING;
  // this is the new process to run
  current_process=chosen;
  // switch to its address space
  write_cr3(chosen->cr3);

  // release mutex
  lock_release_sti(&scheduler_mutex,hwi);
}


process_t *create_process(image_t *exec_image)
{
  process_t *proc;
  dword_t esp;
  uregs_t *stack_regs;
  struct exec *aout_header;
  int hwi;

  // acquire mutex
  hwi=lock_acquire_cli(&scheduler_mutex);

  // Allocate new process descriptor
  proc=(process_t *)malloc(sizeof(process_t));
  
  // get a new PID
  proc->pid=latest_pid++;
  // the process has been just created...
  // it needs to be admitted
  proc->state=PS_CREATED;
  proc->time_slice=TIME_SLICE;
  proc->priority=0;
  
  // real-time parameters (in milli-seconds)
  // they will be filled-in by the system-call
  proc->phase = 0;
  proc->deadline = 0;
  proc->period = 0;
  proc->exec_time = 0;

  proc->absolute_deadline = 0;
  proc->exec_remaining = 0;

  // job type
  proc->type = TYPE_NORMAL;

  // initialize a.out header pointer
  //aout_header = (exec_image != NULL) ? (struct exec *)(exec_image->data) : NULL;
  
  if (exec_image != NULL)
    aout_header=(struct exec *)(exec_image->data);
  else
    aout_header=NULL;

  // initialize image file pointer and a.out header pointer
  // or NULL if not available
  proc->image_file=exec_image;
  proc->aout_header=aout_header;

  // initialize general purpose registers
  process_initialize_registers(proc);

  // initialize a temporary variable esp
  esp=proc->kstack_top;
  stack_regs=(uregs_t *)esp;

  //printf("Creating process with KSTACK esp = %x\n",esp);
  //printf("Address of process struct is = %x\n",proc);
  //printf("Address of process struct kstack is = %x\n",&(proc->kstack[0]));
  //halt();


  // if we have an executable image, use it
  if (exec_image != NULL)
    {
      // copy process name
      strcpy(proc->name,exec_image->name);
      // initialize memory map
      proc->memory_map=initialize_memory_map(exec_image);
      // allocate a new address space
      proc->cr3=new_address_space();
      // user stack top 
      proc->ustack_top = USER_STACK_TOP;
      // assume process is swappable (user process)
      // if it's a system process, this parameter will be set to 0 somewhere else
      proc->swappable=1;

      stack_regs->user_esp=USER_STACK_TOP;
      stack_regs->user_ss=GDT_DESC_UDATA | 3;
      // 3 means DPL 3 -- i.e. User Mode
      stack_regs->cs=GDT_DESC_UCODE | 3;
      stack_regs->ds=GDT_DESC_UDATA | 3;
      stack_regs->es=GDT_DESC_UDATA | 3;
      stack_regs->fs=GDT_DESC_UDATA | 3;
      stack_regs->gs=GDT_DESC_UDATA | 3;
      
      // set entry point
      stack_regs->eip=aout_get_entry_point(aout_header);
    }
  else
    {
      strcpy(proc->name,"unnamed");
      proc->memory_map=NULL;
      // this will be set somewhere else
      proc->cr3=0;
      // process is not swappable, because we have nowhere to load it from
      // because exec_image is NULL
      proc->swappable=0;

      // user kernel stack for user stack as well
      stack_regs->user_esp=esp;
      // stay in kernel mode (we have no user address space)
      stack_regs->user_ss=GDT_DESC_DATA;

      stack_regs->cs=GDT_DESC_CODE;
      stack_regs->ds=GDT_DESC_DATA;
      stack_regs->es=GDT_DESC_DATA;
      stack_regs->fs=GDT_DESC_DATA;
      stack_regs->gs=GDT_DESC_DATA;

      // this will be set somewhere else
      stack_regs->eip=0;    
    }

  // add this process to the "pending for acceptance" list
  llist_add_last(ap_process_list,proc);

  printf("Process \"%s\" has been created.\n",proc->name);

  // release mutex
  lock_release_sti(&scheduler_mutex,hwi);

  return proc;
}


// this works just for user-mode processes
// it's used to refire a periodic process
// do not use it for a kernel process... they can't be periodic
void process_initialize_registers(process_t *proc)
{
  dword_t esp;
  uregs_t *stack_regs;

  esp=(dword_t)&(proc->kstack[PROC_STACK_SIZE]);
  esp-=sizeof(uregs_t);
  proc->kstack_top=esp;
  stack_regs=(uregs_t *)esp;

  stack_regs->ebp=esp;
  stack_regs->esp=esp;

  stack_regs->eax=0x0;
  stack_regs->ebx=0x0;
  stack_regs->ecx=0x0;
  stack_regs->edx=0x0;
  stack_regs->esi=0x0;
  stack_regs->edi=0x0;

  // interrupts are enabled
  stack_regs->eflags=0x202;
  
  // if it's a user process (which has an image)
  if (proc->aout_header != NULL)
    {
      // initialize user-mode stack
      proc->ustack_top = USER_STACK_TOP;
      stack_regs->user_esp=USER_STACK_TOP;
      
      // set entry point  
      stack_regs->eip=aout_get_entry_point(proc->aout_header);
    }
}


process_memory_t *initialize_memory_map(image_t *exec_image)
{
  process_memory_t *map;
  struct exec *aout_header;

  // initialize a.out header pointer
  aout_header=(struct exec *)(exec_image->data);

  // create memory map
  map=(process_memory_t *)malloc(sizeof(process_memory_t));

  // Allocate each section
  
  // Section: text
  map->text=(memory_section_t *)malloc(sizeof(memory_section_t));
  map->text->start = aout_get_text_base(aout_header);
  map->text->length = aout_header->a_text;
  map->text->pages = calculate_num_pages(aout_header->a_text);
  map->text->addresses = llist_make_list();

  // Section: data
  map->data=(memory_section_t *)malloc(sizeof(memory_section_t));
  map->data->start = aout_get_data_base(aout_header);
  map->data->length = aout_header->a_data;
  map->data->pages = calculate_num_pages(aout_header->a_data);
  map->data->addresses = llist_make_list();

  // Section: bss
  map->bss=(memory_section_t *)malloc(sizeof(memory_section_t));
  map->bss->start = aout_get_bss_base(aout_header);
  map->bss->length = aout_header->a_bss;
  map->bss->pages = calculate_num_pages(aout_header->a_bss);
  map->bss->addresses = llist_make_list();

  // Section: heap
  map->heap=(memory_section_t *)malloc(sizeof(memory_section_t));
  map->heap->start = map->bss->start + map->bss->pages * 4096;   //end of BSS and page aligned
  map->heap->length = 0;
  map->heap->pages = 0;
  map->heap->addresses = llist_make_list();

  // Section: stack
  map->stack=(memory_section_t *)malloc(sizeof(memory_section_t));
  map->stack->start = USER_STACK_TOP;
  map->stack->length = 0;
  map->stack->pages = 0;
  map->stack->addresses = llist_make_list();

  return map;  
}


// deallocate all of process memory
void release_memory_map(process_memory_t *map)
{

}

// deallocate all of process memory
void destroy_process(process_t *proc)
{
  // set state to dead
  proc->state=PS_DEAD;

  // if it's the current process we need to change it... so
  if (proc==current_process)
    schedule();

  /*
  // remove the process from the list

  // at some point call this:
  release_memory_map(proc->memory_map);
  free(proc->memory_map);

  // release the memory associated to it
  free(proc);
  */

  // Warning: the caller needs to set *proc=NULL if appropriate;
}


int calculate_num_pages(dword_t byte_size)
{
  int pages;

  pages = byte_size / 4096;
  if (byte_size % 4096)
    pages++;

  return pages;
}


void print_image_list()
{
  ptr_iterator_t iter;
  image_t *img;
  struct exec *aout_header;
  dword_t tb,db,bb,ep;

  iter=llist_get_nodes(process_image_list);

  while (llist_iter_has_next(iter))
    {
      img=llist_iter_next(iter);
      aout_header=(struct exec *)(img->data);

      printf("File: %s -- size = %d\n",img->name,img->size);

      tb=aout_get_text_base(aout_header);
      db=aout_get_data_base(aout_header);
      bb=aout_get_bss_base(aout_header);
      ep=aout_get_entry_point(aout_header);

      printf("  A.OUT (load): text=%x data=%x bss=%x entry=%x\n",tb,db,bb,ep);
      printf("        (size): text=%x data=%x bss=%x\n",aout_header->a_text,aout_header->a_data,aout_header->a_bss);
    }

  free(iter);
}

int aout_loader(process_t *proc,dword_t address)
{
  struct exec *aout_header;
  image_t *image_file;
  byte_t *file_data;
  dword_t palign_addr,frel_addr;
  process_memory_t *mm;
  memory_section_t *ms;
  byte_t memory_access_flags;
  dword_t *ptr_addr;
  int load_from_file;

  if (proc==NULL)
    return 3;

  aout_header=proc->aout_header;
  image_file=proc->image_file;

  // process is unswappable... we have nowehre to load from
  if (aout_header == NULL)
    return 4;

  // point to file data
  file_data=image_file->data;
  
  // get process' memory map
  mm=proc->memory_map;

  // calculate page-aligned address
  palign_addr=address & 0xFFFFF000;

  //  printf("A.OUT Loader: asked to load page %x (%x).\n",palign_addr,address);
  //printf("              for process: \"%s\"\n",proc->name);

  // calculate file-relative address (skip one page and header)
  frel_addr=palign_addr - 0x1000;
  
  // assume we don't have to load from file
  load_from_file=0;

  // check if address should be allocated
  // determine what section it belongs to
  
  // let's see where the fault happened
  if (address < 0x1000)
    {
      return 1;  // null pointer exception
    }
  // text?
  else if (address < (mm->text->start + mm->text->length))
    {
      // point to page
      file_data += frel_addr;
      // code is read-only memory
      //memory_access_flags = VMM_PG_USER | VMM_PG_WRITE;
      memory_access_flags = VMM_PG_USER;
      // point to "text" memory list
      ms = mm->text;
      // we have to load from file (text section)
      load_from_file=1;
      //printf("Loading a TEXT section.\n");
    }
  // data?
  else if (address < (mm->data->start + mm->data->length))
    {
      file_data += frel_addr;
      //printf("Loading a DATA file offset: %x.\n",file_data-image_file->data);
      // read-write memory
      memory_access_flags = VMM_PG_USER | VMM_PG_WRITE;
      // point to "data" memory list
      ms = mm->data;
      // we have to load from file (data section)
      load_from_file=1;
      //printf("Loading a DATA section.\n");
    }
  // bss?
  else if (address < (mm->bss->start + mm->bss->length))
    {
      // point to page
      // skip text and data
      frel_addr -= mm->text->length + mm->data->length;
      file_data += aout_header->a_text + aout_header->a_data + frel_addr;

      // read-write memory
      memory_access_flags = VMM_PG_USER | VMM_PG_WRITE;
      // point to "bss" memory list
      ms = mm->bss;
      //printf("Loading a BSS section.\n");
    }
  // heap?
  //else if (address < (mm->heap->start + mm->heap->length))
  // heap may not overlap with stack
  //else if ((address >= mm->heap->start) && (address < mm->heap->start + mm->heap->length) && (address < (mm->stack->start - mm->stack->length)))  

  /*
  else if ((address >= mm->heap->start) && (address <= (mm->heap->start + mm->heap->length)) && (address < (mm->stack->start - mm->stack->length)))  
    {
      // point to page
      // skip text and data and bss
      frel_addr -= mm->text->length + mm->data->length + mm->bss->length;
      file_data += aout_header->a_text + aout_header->a_data + aout_header->a_bss + frel_addr;

      // read-write memory
      memory_access_flags = VMM_PG_USER | VMM_PG_WRITE;
      // point to "heap" memory list
      ms = mm->heap;
      // add one page to this dynamic area
      ms->length += 4096;
      ms->pages++;
      printf("Loading a HEAP section.\n");
      printf("Heap is at: %x\n",mm->heap->start);
    }
  */
  // stack?
  // stack grows downwards
  // stack may not overlap with heap
  else if ((address > (mm->heap->start + mm->heap->length)) && (address <= mm->stack->start))  
    {
      // read-write memory
      memory_access_flags = VMM_PG_USER | VMM_PG_WRITE;
      // point to "stack" memory list
      ms = mm->stack;

      // add one page to this dynamic area
      ms->length += 4096;
      ms->pages++;
      //printf("Loading a STACK section.\n");
    }
  else
    return 2;  // nope, it's an invalid memory access, page fault, kill process.

  
  // if we're here, then file_data points to 
  // the appropriate file location to load in memory
  // thus allocate one page and load it from file_data
  // Warning: can be improved to check how many bytes to load exactly
  // but for now, just copying one page will do
  
  // allocate one page and map it at palign_addr
  // Warning: need to check for out of memory... but no problem for now
  allocate_page(palign_addr,memory_access_flags);

  // mark allocation on list (value needs to be a pointer to it)
  ptr_addr=(dword_t *)malloc(sizeof(dword_t));
  *ptr_addr=palign_addr;
  llist_add_last(ms->addresses,ptr_addr);
  
  // copy data from file_data to user memory
  // if appropriate (text or data section)
  if (load_from_file)
    {
      //printf("Loading address %x from file offset: %x.\n",palign_addr,file_data-image_file->data);
      memcpy(palign_addr,file_data,4096);
    }

  /*
  printf("Copied (%x Bytes) from %x to %x\n",4096,file_data,palign_addr);
  memory_dump(palign_addr,100);
  printf("Halting in AOUT.");
  halt();
  */

  // everything's ok
  return 0;
}

image_t *find_image(char *name)
{
  ptr_iterator_t iter;
  image_t *img,*found;

  found=NULL;
  iter=llist_get_nodes(process_image_list);

  while (llist_iter_has_next(iter))
    {
      img=llist_iter_next(iter);

      if (!strcmp(img->name,name))
	{
	  found=img;
	  break;
	}
    }  

  free(iter);
  
  return found;
}

// Warning: this function is called by the PIT interrupt
// no need to lock the process list, cause interrupts are disabled
int release_periodic(qword_t current_time)
{
  ptr_node_t pnode;
  process_t *p;
  dword_t ct;
  int released;

  // hack... this works only for 25 days...
  ct=(dword_t)(current_time & (qword_t)0x0FFFFFFFFL);

  released=0;
  
  for (pnode=rt_process_list->header->next; pnode!=rt_process_list->header; pnode=pnode->next)
    {
      p=(process_t *)(pnode->data);

      // release periodic jobs if applicable
      // or remove non-periodic dead ones
      if (p->state == PS_TERMINATED)
	{
	  //releaseable = (((ct - p->phase) % (p->period)) == 0);
	  //if ((p->type == TYPE_PERIODIC) && RELEASEABLE(current_time,p))
	  if ((p->type == TYPE_PERIODIC) && (((ct - p->phase) % (p->period)) == 0))
	    {
	      // refill its execution time
	      p->exec_remaining = p->exec_time;
	      // reset its absolute deadline
	      p->absolute_deadline = current_time + p->deadline;
	      process_initialize_registers(p);
	      p->state = PS_READY;
	      released=1;
	    }

	  // TO DO:
	  // if p->type == TYPE_NORMAL then remove it and delete its address space
	}
    }  
  return released;
}

int acceptancy_test(process_t *p)
{
  double density;
  int accepted;
  int hwi;

  // acquire scheduler mutex  
  hwi=lock_acquire_cli(&scheduler_mutex);

  // remove it from the pending list
  llist_remove_reference(ap_process_list,p);

  if (p->type == TYPE_PERIODIC)
    {
      density = (double)(p->exec_time) / (double)(MIN(p->deadline,p->period));
      
      if ((load_factor + density) < 1)
	{
	  // acceptable
	  load_factor += density;
	  // add it to the Real-Time list
	  llist_add_last(rt_process_list,p);
	  accepted = 1;
	}
      else
	{
	  printf("Can't accept process because load factor would become: %d\n",(dword_t)((load_factor+density)*100));
	  accepted = 0;  // not acceptable
	}
    }
  else  // it's a normal job, accept it anyway (they can only run in idle time)
    {
      // add it to the Round-Robin list
      llist_add_first(rr_process_list,p);
      accepted = 1;
    }

  // release scheduler mutex
  lock_release_sti(&scheduler_mutex,hwi);

  printf("Current load factor: %d\n",(dword_t)(load_factor*100));

  return accepted;
}
